import { create } from "zustand";
import { persist } from "zustand/middleware";

type AIMessageType = "C" | "R";

type Message = {
  from: "human" | "ai";
  id: number;
  data: string;
};

const COUNTER = { value: 0 };

const incrAndGet = () => ++COUNTER.value;

interface ChatState {
  messages: Message[];
  addHumanMessage: (message: string) => void;
  appendAiMessageToken: (
    id: string,
    type: AIMessageType,
    token: string,
  ) => void;
}

export const useChatMessages = create<ChatState>()(
  persist(
    (set) => ({
      messages: [],
      addHumanMessage: (text) =>
        set(({ messages }) => ({
          messages: [
            ...messages,
            { from: "human", id: incrAndGet(), data: text },
          ],
        })),
      appendAiMessageToken: (id, type, token) => {
        set(({ messages }) => {
          const newMessages = [...messages];
          if (
            newMessages.length === 0 ||
            newMessages[newMessages.length - 1].from === "human"
          ) {
            newMessages.push({ from: "ai", id: incrAndGet(), data: "" });
          }

          newMessages[newMessages.length - 1].data += token;

          return { messages: newMessages };
        });
      },
    }),
    {
      name: "chat-messages",
    },
  ),
);
